"""
SimpleAPI ORM 查询接口

提供简化的查询接口，完全封装 Tortoise ORM 的查询功能。
用户只需要使用这些查询接口，无需了解底层的 Tortoise 查询。
"""

from typing import Any, Dict, List, Optional, Type, Union
from tortoise.queryset import QuerySet as TortoiseQuerySet
from tortoise.queryset import Q as TortoiseQ
from tortoise.functions import Count, Sum, Avg, Max, Min
from tortoise.expressions import F as TortoiseF


class Q:
    """查询条件构建器"""
    
    def __init__(self, **kwargs):
        self._q = TortoiseQ(**kwargs)
    
    def __and__(self, other: 'Q') -> 'Q':
        """AND 操作"""
        result = Q()
        result._q = self._q & other._q
        return result
    
    def __or__(self, other: 'Q') -> 'Q':
        """OR 操作"""
        result = Q()
        result._q = self._q | other._q
        return result
    
    def __invert__(self) -> 'Q':
        """NOT 操作"""
        result = Q()
        result._q = ~self._q
        return result
    
    @classmethod
    def AND(cls, *conditions: 'Q') -> 'Q':
        """多个条件 AND"""
        if not conditions:
            return cls()
        result = conditions[0]
        for condition in conditions[1:]:
            result = result & condition
        return result
    
    @classmethod
    def OR(cls, *conditions: 'Q') -> 'Q':
        """多个条件 OR"""
        if not conditions:
            return cls()
        result = conditions[0]
        for condition in conditions[1:]:
            result = result | condition
        return result
    
    @classmethod
    def NOT(cls, condition: 'Q') -> 'Q':
        """NOT 条件"""
        return ~condition


class F:
    """字段表达式"""
    
    def __init__(self, field_name: str):
        self._f = TortoiseF(field_name)
    
    def __add__(self, other):
        """加法"""
        result = F("")
        result._f = self._f + other
        return result
    
    def __sub__(self, other):
        """减法"""
        result = F("")
        result._f = self._f - other
        return result
    
    def __mul__(self, other):
        """乘法"""
        result = F("")
        result._f = self._f * other
        return result
    
    def __truediv__(self, other):
        """除法"""
        result = F("")
        result._f = self._f / other
        return result


class QueryBuilder:
    """查询构建器"""
    
    def __init__(self, model_class: Type):
        self.model_class = model_class
        self.queryset = model_class.all()
        
    def filter(self, *args, **kwargs) -> 'QueryBuilder':
        """
        添加过滤条件
        
        Args:
            *args: Q 对象
            **kwargs: 字段过滤条件
            
        Returns:
            查询构建器
        """
        # 处理 Q 对象
        q_objects = []
        for arg in args:
            if isinstance(arg, Q):
                q_objects.append(arg._q)
            else:
                q_objects.append(arg)
        
        if q_objects:
            self.queryset = self.queryset.filter(*q_objects, **kwargs)
        else:
            self.queryset = self.queryset.filter(**kwargs)
        return self
        
    def exclude(self, *args, **kwargs) -> 'QueryBuilder':
        """
        添加排除条件
        
        Args:
            *args: Q 对象
            **kwargs: 字段排除条件
            
        Returns:
            查询构建器
        """
        # 处理 Q 对象
        q_objects = []
        for arg in args:
            if isinstance(arg, Q):
                q_objects.append(arg._q)
            else:
                q_objects.append(arg)
        
        if q_objects:
            self.queryset = self.queryset.exclude(*q_objects, **kwargs)
        else:
            self.queryset = self.queryset.exclude(**kwargs)
        return self
        
    def order_by(self, *fields) -> 'QueryBuilder':
        """
        添加排序
        
        Args:
            *fields: 排序字段，前缀 - 表示降序
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.order_by(*fields)
        return self
        
    def limit(self, limit: int) -> 'QueryBuilder':
        """
        限制数量
        
        Args:
            limit: 限制数量
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.limit(limit)
        return self
        
    def offset(self, offset: int) -> 'QueryBuilder':
        """
        设置偏移量
        
        Args:
            offset: 偏移量
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.offset(offset)
        return self
        
    def select_related(self, *fields) -> 'QueryBuilder':
        """
        预加载关联对象（一对一、多对一）
        
        Args:
            *fields: 关联字段名
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.select_related(*fields)
        return self
        
    def prefetch_related(self, *fields) -> 'QueryBuilder':
        """
        预取关联对象（一对多、多对多）
        
        Args:
            *fields: 关联字段名
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.prefetch_related(*fields)
        return self
        
    def distinct(self) -> 'QueryBuilder':
        """
        去重
        
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.distinct()
        return self
        
    def values(self, *fields) -> 'QueryBuilder':
        """
        只返回指定字段的值
        
        Args:
            *fields: 字段名
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.values(*fields)
        return self
        
    def values_list(self, *fields, flat: bool = False) -> 'QueryBuilder':
        """
        返回字段值列表
        
        Args:
            *fields: 字段名
            flat: 是否扁平化（仅当只有一个字段时有效）
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.values_list(*fields, flat=flat)
        return self
        
    def annotate(self, **kwargs) -> 'QueryBuilder':
        """
        添加聚合字段
        
        Args:
            **kwargs: 聚合字段定义
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.annotate(**kwargs)
        return self
        
    def group_by(self, *fields) -> 'QueryBuilder':
        """
        分组
        
        Args:
            *fields: 分组字段
            
        Returns:
            查询构建器
        """
        self.queryset = self.queryset.group_by(*fields)
        return self
        
    # 执行查询的方法
    async def all(self) -> List:
        """获取所有结果"""
        return await self.queryset
        
    async def first(self) -> Optional:
        """获取第一个结果"""
        return await self.queryset.first()
        
    async def get(self):
        """获取单个结果"""
        return await self.queryset.get()
        
    async def count(self) -> int:
        """统计数量"""
        return await self.queryset.count()
        
    async def exists(self) -> bool:
        """检查是否存在"""
        return await self.queryset.exists()
        
    async def delete(self) -> int:
        """删除匹配的记录"""
        return await self.queryset.delete()
        
    async def update(self, **kwargs) -> int:
        """更新匹配的记录"""
        return await self.queryset.update(**kwargs)
        
    async def bulk_create(self, objects: List[Dict], batch_size: int = 1000) -> List:
        """批量创建"""
        instances = [self.model_class(**obj) for obj in objects]
        created = await self.model_class.bulk_create(instances, batch_size=batch_size)
        return created if created is not None else instances
        
    async def bulk_update(self, instances: List, fields: List[str], batch_size: int = 1000) -> None:
        """批量更新"""
        if instances:
            await self.model_class.bulk_update(instances, fields, batch_size=batch_size)


class Aggregation:
    """聚合函数"""
    
    @staticmethod
    def count(field: str = "*") -> Count:
        """计数"""
        return Count(field)
    
    @staticmethod
    def sum(field: str) -> Sum:
        """求和"""
        return Sum(field)
    
    @staticmethod
    def avg(field: str) -> Avg:
        """平均值"""
        return Avg(field)
    
    @staticmethod
    def max(field: str) -> Max:
        """最大值"""
        return Max(field)
    
    @staticmethod
    def min(field: str) -> Min:
        """最小值"""
        return Min(field)


class QuerySet:
    """
    查询集合类
    
    提供更高级的查询接口
    """
    
    def __init__(self, model_class: Type):
        self.model_class = model_class
        
    def filter(self, *args, **kwargs) -> QueryBuilder:
        """创建过滤查询"""
        return QueryBuilder(self.model_class).filter(*args, **kwargs)
        
    def exclude(self, *args, **kwargs) -> QueryBuilder:
        """创建排除查询"""
        return QueryBuilder(self.model_class).exclude(*args, **kwargs)
        
    def all(self) -> QueryBuilder:
        """获取所有记录的查询"""
        return QueryBuilder(self.model_class)
        
    def order_by(self, *fields) -> QueryBuilder:
        """创建排序查询"""
        return QueryBuilder(self.model_class).order_by(*fields)
        
    async def get(self, **kwargs):
        """获取单个记录"""
        return await self.model_class.get(**kwargs)
        
    async def get_or_none(self, **kwargs):
        """获取单个记录或返回 None"""
        try:
            return await self.model_class.get(**kwargs)
        except self.model_class.DoesNotExist:
            return None
            
    async def create(self, **kwargs):
        """创建记录"""
        return await self.model_class.create(**kwargs)
        
    async def get_or_create(self, defaults: Dict = None, **kwargs):
        """获取或创建记录"""
        return await self.model_class.get_or_create(defaults=defaults or {}, **kwargs)
        
    async def update_or_create(self, defaults: Dict = None, **kwargs):
        """更新或创建记录"""
        return await self.model_class.update_or_create(defaults=defaults or {}, **kwargs)


# 便捷的查询函数
def query_condition(**kwargs) -> Q:
    """创建查询条件"""
    return Q(**kwargs)


def field_expression(field_name: str) -> F:
    """创建字段表达式"""
    return F(field_name)


def and_conditions(*conditions: Q) -> Q:
    """AND 多个条件"""
    return Q.AND(*conditions)


def or_conditions(*conditions: Q) -> Q:
    """OR 多个条件"""
    return Q.OR(*conditions)


def not_condition(condition: Q) -> Q:
    """NOT 条件"""
    return Q.NOT(condition) 